#include <task.h>
#include <assert.h>
#include <string.h>
#include <global.h>
#include <kernel.h>
#include <os.h>
#include <interrupt.h>
#include <paging.h>
#include <memory32.h>
#include <assert.h>
#include <syscall.h>
#include <elf.h>
#include <fs.h>
#include <logf.h>

#pragma GCC diagnostic ignored "-Wimplicit-function-declaration"

task_queue_t task_queue;

static mutex_t task_mutex;
static task_t task_table[OS_TASK_MAX_SIZE];

static uint32_t idle_task_stack[1024];
static void idle_task_entry()
{
    while (TRUE)
    {
        HLT;   
    }
}

void init_task_queue()
{
    kernel_memset(task_table, 0, sizeof(task_table));
    init_mutex(&task_mutex);

    uint32_t ud_selector = get_gdt_table_entry();
    set_gdt_table_entry(ud_selector, 0x0, 0xFFFFFFFF,
        SEG_ATTR_P | SEG_ATTR_DPL3 | SEG_NORMAL | SEG_TYPE_DATA | SEG_TYPE_RW | SEG_ATTR_D);
    task_queue.ud_selector = ud_selector;

    uint32_t uc_selector = get_gdt_table_entry();
    set_gdt_table_entry(uc_selector, 0x0, 0xFFFFFFFF,
        SEG_ATTR_P | SEG_ATTR_DPL3 | SEG_NORMAL | SEG_TYPE_CODE | SEG_TYPE_RW | SEG_ATTR_D);
    task_queue.uc_selector = uc_selector;

    init_list(&task_queue.ready_list);
    init_list(&task_queue.task_list);
    init_list(&task_queue.sleep_list);
    // 当前无运行任务
    task_queue.running_task = NULL;
    // 初始化空闲任务
    init_task(&task_queue.idle_task, "idle task", TASK_LEVEL_SYSTEM, (uint32_t)idle_task_entry, (uint32_t)&idle_task_stack[1024]);
    // 启动任务
    start_task(&task_queue.idle_task);
}

extern void default_task_entry();

void init_default_task()
{
    // default task 代码开始结束位置
    extern uint8_t b_default_task[], e_default_task[];
    // 计算需要拷贝的字节数
    uint32_t copy_size = (uint32_t)(e_default_task - b_default_task);
    // 分配空间
    uint32_t alloc_size = PAGE_SIZE * 10;
    ASSERT(copy_size < alloc_size);

    // 初始化任务
    uint32_t default_start = (uint32_t)default_task_entry;
    init_task(&task_queue.default_task, 
        "default task", 
        TASK_LEVEL_USER, 
        default_start, 
        default_start + alloc_size);
    write_tr(task_queue.default_task.selector);
    task_queue.running_task = &task_queue.default_task;

    set_pde(task_queue.default_task.tss.cr3);

    alloc_pages(default_start, alloc_size, PTE_P | PTE_W | PTE_U);
    kernel_memcpy((void *)default_start, (void *)b_default_task, copy_size);

    // 启动任务
    start_task(&task_queue.default_task);
}

task_t *get_default_task()
{
    return &task_queue.default_task;
}

int init_tss(task_t *task, uint32_t flag, uint32_t entry, uint32_t esp)
{
    uint32_t selector = get_gdt_table_entry();
    if (selector == -1) {
        return selector;
    }
    
    tss_t *tss = &task->tss;
    set_gdt_table_entry(selector, (uint32_t)tss, sizeof(tss_t),
        SEG_ATTR_P | SEG_ATTR_DPL0 | SEG_TYPE_TSS);
    
    kernel_memset(tss, 0, sizeof(tss_t));
    
    uint32_t kernel_stack = alloc_page();
    ASSERT(kernel_stack != 0);

    uint32_t uc_selector = KERNEL_CODE_SEG, ud_selector = KERNEL_DATA_SEG;
    if (flag & TASK_LEVEL_USER) {
        uc_selector = task_queue.uc_selector | SEG_ATTR_CPL3;
        ud_selector = task_queue.ud_selector | SEG_ATTR_CPL3;
    }
    
    tss->eip = entry;
    tss->esp = esp;
    tss->esp0 = kernel_stack + PAGE_SIZE;
    tss->ss = ud_selector;
    tss->ss0 = KERNEL_DATA_SEG;
    tss->es = tss->ds = tss->fs = tss->gs = ud_selector;
    tss->cs = uc_selector;
    tss->eflags = EFLAGS_DEFAULT | EFLAGS_IF;
    uint32_t pde = memory32_create_pde();
    if (pde == 0) {
        free_gdt_table_entry(selector);
        return -1;
    }
    tss->cr3 = pde;
    task->selector = selector;
    return 0;
}

void start_task(task_t *task)
{
    // 插入就绪队列
    protect_state_t ps = protect_enter();
    task_set_ready(task);
    protect_exit(ps);
}

static uint32_t task_pid = 0;

int init_task(task_t *task, const char *name, uint32_t flag, uint32_t entry, uint32_t esp)
{
    ASSERT(task != (task_t *)0);
    int ret = init_tss(task, flag, entry, esp);
    if (ret != 0) {
        return ret;
    }
    kernel_strcpy(task->name, name);
    task->state = TASK_CREATED;
    task->ticks = task->slices = TASK_DEFAULT_TICKS;
    task->sleep = 0;
    task->pid = task_pid++;
    task->parent = NULL;
    init_list_node(&task->task_node);
    init_list_node(&task->running_node);
    // 插入任务队列
    list_insert_front(&task_queue.task_list, &task->task_node);
    return 0;
}

void init_simple_task(simple_task_t *task, uint32_t entry, uint32_t esp)
{
    ASSERT(task != (simple_task_t *)0);
    uint32_t *pesp = (uint32_t *)esp;
    if (pesp) {
        *(--pesp) = entry;
        *(--pesp) = 0;
        *(--pesp) = 1;
        *(--pesp) = 2;
        *(--pesp) = 3;
        task->stack = pesp;
    }
}

void task_switch(task_t *from, task_t *to)
{
    far_jump(to->selector, 0);
}

extern void simple_switch(uint32_t **from, uint32_t *to);

void simple_task_switch(simple_task_t *from, simple_task_t *to)
{
    simple_switch(&from->stack, to->stack);
}

void task_set_ready(task_t *task)
{
    // 防止空闲任务进入运行任务队列
    if (task == &task_queue.idle_task) return;

    task->state = TASK_READY;
    list_insert_back(&task_queue.ready_list, &task->running_node);
}

void task_set_block(task_t *task)
{
    // 防止空闲任务进入运行任务队列
    if (task == &task_queue.idle_task) return;

    list_remove(&task_queue.ready_list, &task->running_node);
}

void task_set_sleep(task_t *task, uint32_t ticks)
{
    if (ticks == 0) return;

    task->sleep = ticks;
    task->state = TASK_SLEEP;
    list_insert_back(&task_queue.sleep_list, &task->running_node);
}

void task_set_wakeup(task_t *task)
{
    list_remove(&task_queue.sleep_list, &task->running_node);
}

task_t *get_running_task()
{
    return task_queue.running_task;
}

void task_dispatch()
{
    protect_state_t ps = protect_enter();
    task_t *to = &task_queue.idle_task, *from = get_running_task();
    // 判断运行队列是否为空，为空则运行空闲任务
    if (!list_is_empty(&task_queue.ready_list))
    {
        list_node_t *node = list_get_first(&task_queue.ready_list);
        to = struct_from_field(node, task_t, running_node);
    }

    if (to != from)
    {
        task_queue.running_task = to;
        to->state = TASK_RUNNING;
        task_switch(from, to);
    }
    protect_exit(ps);
}

void task_time_tick()
{
    task_t *task = get_running_task();
    if (-- task->ticks == 0)
    {
        task->ticks = task->slices;
        task_set_block(task);
        task_set_ready(task);
        task_dispatch();
    }

    list_node_t *node = list_get_first(&task_queue.sleep_list);
    while (node)
    {
        task_t *task = struct_from_field(node, task_t, running_node);
        if (-- task->sleep == 0) {
            task_set_wakeup(task);
            task_set_ready(task);
        }
        node = task->running_node.next;
    }
    task_dispatch();
}

void task_sys_yield()
{
    protect_state_t ps = protect_enter();
    if (task_queue.ready_list.size > 1)
    {
        task_t *task = get_running_task();
        task_set_block(task);
        task_set_ready(task);
        task_dispatch();
    }
    protect_exit(ps);
}

void task_sys_sleep(uint32_t ms)
{
    protect_state_t ps = protect_enter();
    task_set_block(task_queue.running_task);
    task_set_sleep(task_queue.running_task, (ms + (OS_TICKS_MS - 1)) / OS_TICKS_MS);
    task_dispatch();
    protect_exit(ps);
}

uint32_t task_sys_getpid()
{
    task_t *task = get_running_task();
    return task->pid;
}

static task_t *alloc_task()
{
    task_t *task = NULL;

    mutex_lock(&task_mutex);
    for (int i = 0; i < OS_TASK_MAX_SIZE; i++)
    {
        task_t *t = &task_table[i];
        if (t->name[0] == 0) {
            task = t;
            break;
        }
    }
    mutex_unlock(&task_mutex);
    
    return task;
}

static void free_task(task_t *task)
{
    mutex_lock(&task_mutex);
    kernel_memset(task->name, 0, TASK_NAME_SIZE);
    mutex_unlock(&task_mutex);
}

int task_sys_fork()
{
    task_t *parent = get_running_task();
    task_t *child = alloc_task();
    if (child == NULL) 
        return -1;
    
    syscall_frame_t *frame = (syscall_frame_t *)(parent->tss.esp0 - sizeof(syscall_frame_t));
    int ret = init_task(child, parent->name, TASK_LEVEL_USER, 
            frame->eip, 
            frame->_esp + sizeof(uint32_t) * SYSCALL_PMC);
    if (ret != 0) {
        task_destory(child);
        return -1;
    }

    tss_t *tss = &child->tss;
    tss->eax = 0;
    tss->edx = frame->edx;
    tss->ecx = frame->ecx;
    tss->ebx = frame->ebx;
    tss->esi = frame->esi;
    tss->edi = frame->edi;
    tss->ebp = frame->ebp;
    tss->cs = frame->cs;
    tss->ds = frame->ds;
    tss->es = frame->es;
    tss->gs = frame->gs;
    tss->fs = frame->fs;
    tss->eflags = frame->eflags;

    child->parent = parent;

    if ((tss->cr3 = memory32_copy_page(parent->tss.cr3)) < 0) {
        task_destory(child);
        return -1;
    }
    // 启动任务
    start_task(child);
    return child->pid;
}

void task_destory(task_t *task)
{
    if (task->selector) {
        free_gdt_table_entry(task->selector);
    }

    if (task->tss.esp0) {
        free_page(task->tss.esp - PAGE_SIZE);
    }

    if (task->tss.cr3) {
        // TODO: 回收页
        memory32_destory_page(task->tss.cr3);
    }

    kernel_memset(task, 0, sizeof(task_t));
}

static int load_elf_phdr(int file, Elf32_Phdr *phdr, uint32_t pde)
{
    int err = memory32_alloc_task_pages(pde, phdr->p_vaddr, phdr->p_memsz, PTE_P | PTE_U | PTE_W);
    if (err < 0) {
        return -1;
    }
    if (fs_sys_lseek(file, phdr->p_offset, 0) < 0) {
        return -1;
    }

    uint32_t vaddr = phdr->p_vaddr;
    uint32_t size = phdr->p_filesz;
    while (size > 0)
    {
        int current_size = (size > PAGE_SIZE) ? PAGE_SIZE : size;
        uint32_t paddr = memory32_get_paddr(pde, vaddr);

        if (fs_sys_read(file, (char *)paddr, current_size) < current_size) {
            return -1;
        }

        size -= current_size;
        vaddr += current_size;
    }
    
    return 0;
}

static uint32_t load_elf_file(task_t *task, const char *name, uint32_t pde)
{
    Elf32_Ehdr elf_hdr;
    Elf32_Phdr elf_phdr;
    int file = fs_sys_open(name, 0);
    if (file < 0) {
        logf("open failed: %s", name);
        return 0;
    }
    int count = fs_sys_read(file, (char *)&elf_hdr, sizeof(Elf32_Ehdr));
    if (count < sizeof(Elf32_Ehdr)) {
        logf("read failed");
        if (file) {
            fs_sys_close(file);
        }
        return 0;
    }

    if (elf_hdr.e_ident[0] != 0x7F || 
    elf_hdr.e_ident[1] != 'E' ||elf_hdr.e_ident[2] != 'L' ||elf_hdr.e_ident[3] != 'F') {
        logf("elf file read failed");
        if (file) {
            fs_sys_close(file);
        }
        return 0;
    }
    uint32_t e_phoff = elf_hdr.e_phoff;
    for (int i = 0; i < elf_hdr.e_phnum; i++) {
        if (fs_sys_lseek(file, e_phoff, 0) < 0) {
            logf("read file failed");
            if (file) {
                fs_sys_close(file);
            }
            return 0;
        }

        count = fs_sys_read(file, (char *)&elf_phdr, sizeof(elf_phdr));
        if (count < sizeof(elf_phdr)) {
            logf("read file failed");
            if (file) {
                fs_sys_close(file);
            }
            return 0;
        }

        if ((elf_phdr.p_type != 1) || (elf_phdr.p_vaddr < MEMORY_TASK_BASE)) {
            continue;
        }
        if (load_elf_phdr(file, &elf_phdr, pde) < 0) {
            logf("load phdr failed");
            if (file) {
                fs_sys_close(file);
            }
            return 0;
        }
    }
    
    fs_sys_close(file);
    return elf_hdr.e_entry;
}

static int copy_args(char *to, uint32_t pde, int argc, char **argv)
{
    task_args_t task_args;
    task_args.argc = argc;
    task_args.argv = (char **)(to + sizeof(task_args_t));

    char *dst_arg = to + sizeof(task_args_t) + sizeof(char *) * argc;
    char **dst_arg_tb = memory32_get_paddr(pde, (uint32_t)(to + sizeof(task_args_t)));
    for (int i = 0; i < argc; i++)
    {
        char *from = argv[i];
        int len = kernel_strlen(from) + 1;
        int err = memory32_copy_page_data((uint32_t)dst_arg, pde, (uint32_t)from, len);
        ASSERT(err >= 0);
        dst_arg_tb[i] = dst_arg;
        dst_arg += len;
    }
    
    memory32_copy_page_data((uint32_t)to, pde, (uint32_t)&task_args, sizeof(task_args));
}

int task_sys_execve(char *name, char *argv[], char *env[])
{
    task_t *task = get_running_task();

    kernel_strncpy(task->name, get_file_name(name), TASK_NAME_SIZE);

    uint32_t old_pde = task->tss.cr3;
    uint32_t pde = memory32_create_pde();
    if (!pde) {
        return -1;
    }
    uint32_t entry = load_elf_file(task, name, pde);
    if (entry == 0) {
        task->tss.cr3 = old_pde;
        set_pde(old_pde);
        memory32_destory_page(old_pde);
        return -1;
    }

    uint32_t stack_top = MEMORY_TASK_STACK_TOP - MEMORY_TASK_ARG_SIZE;
    int err = memory32_alloc_task_pages(pde, MEMORY_TASK_STACK_TOP - MEMORY_TASK_STACK_SIZE, MEMORY_TASK_STACK_SIZE, PTE_P | PTE_U | PTE_W);
    if (err < 0) {
        task->tss.cr3 = old_pde;
        set_pde(old_pde);
        memory32_destory_page(old_pde);
        return -1;
    }

    int argc = strings_count(argv);
    err = copy_args((char *)stack_top, pde, argc, argv);
    if (err < 0) {
        task->tss.cr3 = old_pde;
        set_pde(old_pde);
        memory32_destory_page(old_pde);
        return -1;
    }

    syscall_frame_t *frame = (syscall_frame_t *)(task->tss.esp0 - sizeof(syscall_frame_t));
    frame->eip = entry;
    frame->eax = frame->ebx = frame->ecx = frame->edx = 0;
    frame->edi = frame->esi = frame->ebp = 0;
    frame->eflags = EFLAGS_DEFAULT | EFLAGS_IF;
    frame->_esp = stack_top - sizeof(uint32_t) * SYSCALL_PMC;

    task->tss.cr3 = pde;
    set_pde(pde);
    memory32_destory_page(old_pde);
    return 0;
}